{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "rational-charlotte",
   "metadata": {},
   "source": [
    "# Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "opened-facing",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.autograd as autograd         # computation graph\n",
    "from torch import Tensor                  # tensor node in the computation graph\n",
    "import torch.nn as nn                     # neural networks\n",
    "import torch.optim as optim               # optimizers e.g. gradient descent, ADAM, etc.\n",
    "import time\n",
    "from pyDOE import lhs         #Latin Hypercube Sampling\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker\n",
    "\n",
    "#Set default dtype to float32\n",
    "torch.set_default_dtype(torch.float)\n",
    "\n",
    "#PyTorch random number generator\n",
    "torch.manual_seed(1234)\n",
    "\n",
    "# Random number generators in other libraries\n",
    "np.random.seed(1234)\n",
    "\n",
    "# Device configuration\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "print(device)\n",
    "\n",
    "if device == 'cuda': print(torch.cuda.get_device_name()) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "excessive-twist",
   "metadata": {},
   "source": [
    "# *Data Prep*\n",
    "\n",
    "Training and Testing data is prepared from the solution file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "greater-lemon",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_1 = np.linspace(-1,1,256)  # 256 points between -1 and 1 [256x1]\n",
    "x_2 = np.linspace(1,-1,256)  # 256 points between 1 and -1 [256x1]\n",
    "\n",
    "X, Y = np.meshgrid(x_1,x_2) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "reliable-victim",
   "metadata": {},
   "source": [
    "# Test Data\n",
    "\n",
    "We prepare the test data to compare against the solution produced by the PINN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "unlike-windows",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_u_test = np.hstack((X.flatten(order='F')[:,None], Y.flatten(order='F')[:,None]))\n",
    "\n",
    "# Domain bounds\n",
    "lb = np.array([-1, -1]) #lower bound\n",
    "ub = np.array([1, 1])  #upper bound\n",
    "\n",
    "a_1 = 1 \n",
    "a_2 = 1\n",
    "\n",
    "k = 1\n",
    "\n",
    "usol = np.sin(a_1 * np.pi * X) * np.sin(a_2 * np.pi * Y) #solution chosen for convinience  \n",
    "\n",
    "u_true = usol.flatten('F')[:,None] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "assisted-gnome",
   "metadata": {},
   "source": [
    "# Training Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "continuous-liability",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainingdata(N_u,N_f):\n",
    "    \n",
    "    leftedge_x = np.hstack((X[:,0][:,None], Y[:,0][:,None]))\n",
    "    leftedge_u = usol[:,0][:,None]\n",
    "    \n",
    "    rightedge_x = np.hstack((X[:,-1][:,None], Y[:,-1][:,None]))\n",
    "    rightedge_u = usol[:,-1][:,None]\n",
    "    \n",
    "    topedge_x = np.hstack((X[0,:][:,None], Y[0,:][:,None]))\n",
    "    topedge_u = usol[0,:][:,None]\n",
    "    \n",
    "    bottomedge_x = np.hstack((X[-1,:][:,None], Y[-1,:][:,None]))\n",
    "    bottomedge_u = usol[-1,:][:,None]\n",
    "    \n",
    "    all_X_u_train = np.vstack([leftedge_x, rightedge_x, bottomedge_x, topedge_x])\n",
    "    all_u_train = np.vstack([leftedge_u, rightedge_u, bottomedge_u, topedge_u])  \n",
    "     \n",
    "    #choose random N_u points for training\n",
    "    idx = np.random.choice(all_X_u_train.shape[0], N_u, replace=False) \n",
    "    \n",
    "    X_u_train = all_X_u_train[idx[0:N_u], :] #choose indices from  set 'idx' (x,t)\n",
    "    u_train = all_u_train[idx[0:N_u],:]      #choose corresponding u\n",
    "    \n",
    "    '''Collocation Points'''\n",
    "\n",
    "    # Latin Hypercube sampling for collocation points \n",
    "    # N_f sets of tuples(x,t)\n",
    "    X_f = lb + (ub-lb)*lhs(2,N_f) \n",
    "    X_f_train = np.vstack((X_f, X_u_train)) # append training points to collocation points \n",
    "    \n",
    "    return X_f_train, X_u_train, u_train \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stretch-dating",
   "metadata": {},
   "source": [
    "# Physics Informed Neural Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "yellow-intersection",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sequentialmodel(nn.Module):\n",
    "    \n",
    "    def __init__(self,layers):\n",
    "        super().__init__() #call __init__ from parent class \n",
    "              \n",
    "        'activation function'\n",
    "        self.activation = nn.Tanh()\n",
    "\n",
    "        'loss function'\n",
    "        self.loss_function = nn.MSELoss(reduction ='mean')\n",
    "    \n",
    "        'Initialise neural network as a list using nn.Modulelist'  \n",
    "        self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i+1]) for i in range(len(layers)-1)])\n",
    "        \n",
    "        '''\n",
    "        Alternatively:\n",
    "        \n",
    "        *all layers are callable \n",
    "    \n",
    "        Simple linear Layers\n",
    "        self.fc1 = nn.Linear(2,50)\n",
    "        self.fc2 = nn.Linear(50,50)\n",
    "        self.fc3 = nn.Linear(50,50)\n",
    "        self.fc4 = nn.Linear(50,1)\n",
    "        \n",
    "        '''\n",
    "    \n",
    "        'Xavier Normal Initialization'\n",
    "        # std = gain * sqrt(2/(input_dim+output_dim))\n",
    "        for i in range(len(layers)-1):\n",
    "            \n",
    "            # weights from a normal distribution with \n",
    "            # Recommended gain value for tanh = 5/3?\n",
    "            nn.init.xavier_normal_(self.linears[i].weight.data, gain=1.0)\n",
    "            \n",
    "            # set biases to zero\n",
    "            nn.init.zeros_(self.linears[i].bias.data)\n",
    "            \n",
    "\n",
    "    def forward(self,x):\n",
    "        \n",
    "        if torch.is_tensor(x) != True:         \n",
    "            x = torch.from_numpy(x)                \n",
    "        \n",
    "        u_b = torch.from_numpy(ub).float().to(device)\n",
    "        l_b = torch.from_numpy(lb).float().to(device)\n",
    "                      \n",
    "        #preprocessing input \n",
    "        x = (x - l_b)/(u_b - l_b) #feature scaling\n",
    "        \n",
    "        #convert to float\n",
    "        a = x.float()\n",
    "                        \n",
    "        '''     \n",
    "        Alternatively:\n",
    "        \n",
    "        a = self.activation(self.fc1(a))\n",
    "        a = self.activation(self.fc2(a))\n",
    "        a = self.activation(self.fc3(a))\n",
    "        a = self.fc4(a)\n",
    "        \n",
    "        '''\n",
    "        \n",
    "        for i in range(len(layers)-2):\n",
    "            \n",
    "            z = self.linears[i](a)\n",
    "                        \n",
    "            a = self.activation(z)\n",
    "            \n",
    "        a = self.linears[-1](a)\n",
    "        \n",
    "        return a\n",
    "                        \n",
    "    def loss_BC(self,x,y):\n",
    "                \n",
    "        loss_u = self.loss_function(self.forward(x), y)\n",
    "                \n",
    "        return loss_u\n",
    "    \n",
    "    def loss_PDE(self, x_to_train_f):\n",
    "                \n",
    "        x_1_f = x_to_train_f[:,[0]]\n",
    "        x_2_f = x_to_train_f[:,[1]]\n",
    "                        \n",
    "        g = x_to_train_f.clone()\n",
    "                        \n",
    "        g.requires_grad = True\n",
    "        \n",
    "        u = self.forward(g)\n",
    "                \n",
    "        u_x = autograd.grad(u,g,torch.ones([x_to_train_f.shape[0], 1]).to(device), retain_graph=True, create_graph=True)[0]\n",
    "                                \n",
    "        u_xx = autograd.grad(u_x,g,torch.ones(x_to_train_f.shape).to(device), create_graph=True)[0]\n",
    "                                                            \n",
    "        u_xx_1 = u_xx[:,[0]]\n",
    "        \n",
    "        u_xx_2 = u_xx[:,[1]]\n",
    "                \n",
    "        q = ( -(a_1*np.pi)**2 - (a_2*np.pi)**2 + k**2 ) * torch.sin(a_1*np.pi*x_1_f) * torch.sin(a_2*np.pi*x_2_f)\n",
    "                        \n",
    "        f = u_xx_1 + u_xx_2 + k**2 * u - q \n",
    "        \n",
    "        loss_f = self.loss_function(f,f_hat)\n",
    "                \n",
    "        return loss_f\n",
    "    \n",
    "    def loss(self,x,y,x_to_train_f):\n",
    "\n",
    "        loss_u = self.loss_BC(x,y)\n",
    "        loss_f = self.loss_PDE(x_to_train_f)\n",
    "        \n",
    "        loss_val = loss_u #+ loss_f\n",
    "        \n",
    "        return loss_val\n",
    "     \n",
    "    'callable for optimizer'                                       \n",
    "    def closure(self):\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        loss_val = self.loss(X_u_train, u_train, X_f_train)\n",
    "        \n",
    "        error_vec, _ = PINN.test()\n",
    "        \n",
    "        print(loss,error_vec)\n",
    "        \n",
    "        loss_val.backward()\n",
    "\n",
    "        return loss_val        \n",
    "    \n",
    "    def test(self):\n",
    "                \n",
    "        u_pred = self.forward(X_u_test_tensor)\n",
    "        \n",
    "        error_vec = torch.linalg.norm((u-u_pred),2)/torch.linalg.norm(u,2)        # Relative L2 Norm of the error (Vector)\n",
    "        \n",
    "        u_pred = np.reshape(u_pred.cpu().detach().numpy(),(256,256),order='F') \n",
    "        \n",
    "        return error_vec, u_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "thorough-living",
   "metadata": {
    "colab_type": "text",
    "id": "bOjuHdzAhib-"
   },
   "source": [
    "# *Solution Plot*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "suitable-wagon",
   "metadata": {},
   "outputs": [],
   "source": [
    "def solutionplot(u_pred,X_u_train,u_train):\n",
    "\n",
    "    #Ground truth\n",
    "    fig_1 = plt.figure(1, figsize=(18, 5))\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.pcolor(x_1, x_2, usol, cmap='jet')\n",
    "    plt.colorbar()\n",
    "    plt.xlabel(r'$x_1$', fontsize=18)\n",
    "    plt.ylabel(r'$x_2$', fontsize=18)\n",
    "    plt.title('Ground Truth $u(x_1,x_2)$', fontsize=15)\n",
    "\n",
    "    # Prediction\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.pcolor(x_1, x_2, u_pred, cmap='jet')\n",
    "    plt.colorbar()\n",
    "    plt.xlabel(r'$x_1$', fontsize=18)\n",
    "    plt.ylabel(r'$x_2$', fontsize=18)\n",
    "    plt.title('Predicted $\\hat u(x_1,x_2)$', fontsize=15)\n",
    "\n",
    "    # Error\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.pcolor(x_1, x_2, np.abs(usol - u_pred), cmap='jet')\n",
    "    plt.colorbar()\n",
    "    plt.xlabel(r'$x_1$', fontsize=18)\n",
    "    plt.ylabel(r'$x_2$', fontsize=18)\n",
    "    plt.title(r'Absolute error $|u(x_1,x_2)- \\hat u(x_1,x_2)|$', fontsize=15)\n",
    "    plt.tight_layout()\n",
    "\n",
    "    plt.savefig('Helmholtz_non_stiff.png', dpi = 500, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "engaged-supply",
   "metadata": {},
   "source": [
    "# Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "strange-opening",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequentialmodel(\n",
      "  (activation): Tanh()\n",
      "  (loss_function): MSELoss()\n",
      "  (linears): ModuleList(\n",
      "    (0): Linear(in_features=2, out_features=50, bias=True)\n",
      "    (1): Linear(in_features=50, out_features=50, bias=True)\n",
      "    (2): Linear(in_features=50, out_features=50, bias=True)\n",
      "    (3): Linear(in_features=50, out_features=1, bias=True)\n",
      "  )\n",
      ")\n",
      "tensor(0.1682, grad_fn=<MseLossBackward>) tensor(1.1221, grad_fn=<DivBackward0>)\n",
      "tensor(1.9169e-05, grad_fn=<MseLossBackward>) tensor(1.0015, grad_fn=<DivBackward0>)\n",
      "tensor(1.4385e-05, grad_fn=<MseLossBackward>) tensor(1.0012, grad_fn=<DivBackward0>)\n",
      "tensor(1.1532e-05, grad_fn=<MseLossBackward>) tensor(1.0009, grad_fn=<DivBackward0>)\n",
      "tensor(9.0060e-06, grad_fn=<MseLossBackward>) tensor(1.0006, grad_fn=<DivBackward0>)\n",
      "tensor(6.9165e-06, grad_fn=<MseLossBackward>) tensor(1.0004, grad_fn=<DivBackward0>)\n",
      "tensor(5.2457e-06, grad_fn=<MseLossBackward>) tensor(1.0002, grad_fn=<DivBackward0>)\n",
      "tensor(3.9330e-06, grad_fn=<MseLossBackward>) tensor(1.0001, grad_fn=<DivBackward0>)\n",
      "tensor(2.9159e-06, grad_fn=<MseLossBackward>) tensor(1.0000, grad_fn=<DivBackward0>)\n",
      "tensor(2.1427e-06, grad_fn=<MseLossBackward>) tensor(0.9999, grad_fn=<DivBackward0>)\n",
      "Training time: 25.59\n",
      "Test Error: 0.99992\n"
     ]
    }
   ],
   "source": [
    "N_u = 400 #Total number of data points for 'u'\n",
    "N_f = 10000 #Total number of collocation points \n",
    "\n",
    "# Training data\n",
    "X_f_train_np_array, X_u_train_np_array, u_train_np_array = trainingdata(N_u,N_f)\n",
    "\n",
    "'Convert to tensor and send to GPU'\n",
    "X_f_train = torch.from_numpy(X_f_train_np_array).float().to(device)\n",
    "X_u_train = torch.from_numpy(X_u_train_np_array).float().to(device)\n",
    "u_train = torch.from_numpy(u_train_np_array).float().to(device)\n",
    "X_u_test_tensor = torch.from_numpy(X_u_test).float().to(device)\n",
    "u = torch.from_numpy(u_true).float().to(device)\n",
    "f_hat = torch.zeros(X_f_train.shape[0],1).to(device)\n",
    "\n",
    "layers = np.array([2, 50, 50, 50, 1]) #3 hidden layers\n",
    "\n",
    "PINN = Sequentialmodel(layers)\n",
    "       \n",
    "PINN.to(device)\n",
    "\n",
    "'Neural Network Summary'\n",
    "\n",
    "print(PINN)\n",
    "\n",
    "params = list(PINN.parameters())\n",
    "\n",
    "'''Optimization'''\n",
    "\n",
    "'L-BFGS Optimizer'\n",
    "\n",
    "# optimizer = torch.optim.LBFGS(PINN.parameters(), lr=0.1, \n",
    "#                               max_iter = 1000, \n",
    "#                               max_eval = None, \n",
    "#                               tolerance_grad = 1e-06, \n",
    "#                               tolerance_change = 1e-09, \n",
    "#                               history_size = 100, \n",
    "#                               line_search_fn = 'strong_wolfe')\n",
    "\n",
    "\n",
    "# start_time = time.time()\n",
    "\n",
    "# optimizer.zero_grad()     # zeroes the gradient buffers of all parameters\n",
    "\n",
    "# optimizer.step(PINN.closure)\n",
    "\n",
    "\n",
    "'Adam Optimizer'\n",
    "\n",
    "optimizer = optim.Adam(PINN.parameters(), lr=0.001,betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)\n",
    "\n",
    "max_iter = 1000\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "for i in range(max_iter):\n",
    "\n",
    "    loss = PINN.loss(X_u_train, u_train, X_f_train)\n",
    "           \n",
    "    optimizer.zero_grad()     # zeroes the gradient buffers of all parameters\n",
    "    \n",
    "    loss.backward() #backprop\n",
    "\n",
    "    optimizer.step()\n",
    "    \n",
    "    if i % (max_iter/10) == 0:\n",
    "\n",
    "        error_vec, _ = PINN.test()\n",
    "\n",
    "        print(loss,error_vec)\n",
    "    \n",
    "    \n",
    "elapsed = time.time() - start_time                \n",
    "print('Training time: %.2f' % (elapsed))\n",
    "\n",
    "\n",
    "''' Model Accuracy ''' \n",
    "error_vec, u_pred = PINN.test()\n",
    "\n",
    "print('Test Error: %.5f'  % (error_vec))\n",
    "\n",
    "\n",
    "''' Solution Plot '''\n",
    "# solutionplot(u_pred,X_u_train,u_train)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
